//
// reference:
//  https://developer.bitcoin.org/devguide/block_chain.html
//  https://developer.bitcoin.org/reference/block_chain.html
//  https://happypeter.github.io/binfo/nonce
//

use num_bigint::BigUint;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tracing::{error, warn};

use crate::{
    BlockChainHash, BlockId, TransactionId,
    block::{self, BlockError},
    constants::{BLOCK_CHIAN_VERSION, MAX_BLOCK_SIZE, NULL_BLOCK_ID, PROOF_OF_WORK_LIMIT},
    transaction::Transaction,
};

//
// https://developer.bitcoin.org/reference/block_chain.html
//
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BlockHeader {
    version: u32,
    prev_block_header_hash: BlockId,
    merkle_root_hash: BlockChainHash,
    timestamp: u32,
    nbits: u32,
    nonce: u32,
}

impl BlockHeader {
    pub fn new(
        version: u32,
        prev_block_hash: &BlockId,
        merkle_root_hash: BlockChainHash,
        timestamp: u32,
        nbits: u32,
        nonce: u32,
    ) -> Self {
        Self {
            version,
            prev_block_header_hash: *prev_block_hash,
            merkle_root_hash,
            timestamp,
            nbits,
            nonce,
        }
    }

    pub fn new_null() -> Self {
        Self {
            version: BLOCK_CHIAN_VERSION,
            prev_block_header_hash: BlockId::new(NULL_BLOCK_ID),
            merkle_root_hash: BlockChainHash::new([0; 32]),
            timestamp: 0,
            nbits: 0,
            nonce: 0,
        }
    }

    pub fn new_with_nbits_prev_block_id(nbits: u32, prev_block_header_hash: &BlockId) -> Self {
        Self {
            version: BLOCK_CHIAN_VERSION,
            prev_block_header_hash: *prev_block_header_hash, //BlockId::new(NULL_BLOCK_ID),
            merkle_root_hash: BlockChainHash::new([0; 32]),
            timestamp: 0,
            nbits,
            nonce: 0,
        }
    }

    fn check(&self) -> Result<(), BlockError> {
        if self.prev_block_header_hash == BlockId::new(NULL_BLOCK_ID) {
            return Err(BlockError::NoPrevBlock);
        }
        Ok(())
    }
    fn dump(&self) {
        println!(
            r#"  BlockHeader:
    version: {:#x}
    prev_block_hash:  {}
    merkle_root_hash: {}
    timestamp: {}
    nbits: {:#x}
    nonce: {:#x}"#,
            self.version,
            self.prev_block_header_hash.to_hex_string(),
            self.merkle_root_hash.to_hex_string(),
            self.timestamp,
            self.nbits,
            self.nonce
        );
    }
    fn set_merkle_root(&mut self, merkle_root_hash: BlockChainHash) {
        self.merkle_root_hash = merkle_root_hash;
    }

    fn get_merkle_root(&self) -> BlockChainHash {
        self.merkle_root_hash
    }

    fn calc_block_id(&self) -> BlockId {
        let bytes = bincode::serialize(self).unwrap();
        let hash = Sha256::digest(&bytes);
        let hash_bytes: [u8; 32] = hash.into();
        BlockId::new(hash_bytes)
    }

    fn find_and_set_valid_nonce(&mut self, target: &BlockChainHash) -> BlockId {
        let mut block_id = BlockId::new([0; 32]);
        while self.nonce < u32::MAX {
            block_id = self.calc_block_id();
            if block_id < *target {
                return block_id;
            }
            self.nonce += 1;
        }

        block_id
    }

    fn get_nbits(&self) -> u32 {
        self.nbits
    }

    // A 256-bit integer (BigUint) representing the target difficulty.
    pub fn nbits_to_big_uint(&self) -> BigUint {
        block::nbits_to_target_threshold(self.nbits).unwrap()
    }

    pub fn nbits_to_hash(&self) -> BlockChainHash {
        block::nbits_to_hash(self.nbits)
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
struct BlockPayload {
    header: BlockHeader,
    transactions: Vec<Transaction>,
}

impl BlockPayload {
    fn new(header: BlockHeader) -> Self {
        Self {
            header,
            transactions: vec![],
        }
    }

    fn add_transaction(&mut self, tx: Transaction) {
        self.transactions.push(tx);
    }

    fn delete_transaction(&mut self, tx_id: &TransactionId) -> usize {
        let mut seriailze_len = 0;
        self.transactions
            .iter()
            .position(|tx| &tx.calc_tx_id() == tx_id)
            .map(|i| {
                seriailze_len += self.transactions[i].serialize().len();
                self.transactions.remove(i);
                seriailze_len
            })
            .unwrap_or(0)
    }

    fn get_transactions(&self) -> &Vec<Transaction> {
        &self.transactions
    }

    fn dump(&self) {
        self.header.dump();

        print!("  Transactions:");
        for tx in &self.transactions {
            println!("\n    {:#?}", tx);
        }
    }

    fn check(&self) -> Result<(), BlockError> {
        if self.transactions.is_empty() {
            return Err(BlockError::NoTransaction);
        }

        if self.get_serialize_size() > MAX_BLOCK_SIZE {
            return Err(BlockError::BlockSizeTooLarge);
        }

        if !self.transactions[0].is_coinbase() {
            return Err(BlockError::FirstTransactionNotCoinbase);
        }

        for tx in &self.transactions[1..] {
            if tx.is_coinbase() {
                return Err(BlockError::CoinbaseTransactionInMiddle);
            }
        }

        for tx in &self.transactions {
            tx.preliminary_check()?;
        }

        let limit = BigUint::from_bytes_be(&PROOF_OF_WORK_LIMIT);
        let nbits = self.header.get_nbits();
        let target = block::nbits_to_target_threshold(nbits).unwrap();
        if target > limit {
            warn!(
                "NBits below minimum work, nbits: {:#x}, target: {:064x}, limit: {:064x}",
                nbits, target, limit
            );
            return Err(BlockError::NBitsBelowMinimumWork(nbits));
        }

        let block_hash = self.header.calc_block_id();
        let block_hash = BigUint::from_bytes_be(block_hash.as_bytes_ref());
        if block_hash > target {
            error!(
                "Hash does not match nbits. block hash: {:#066x}, nbits: {:#x}, target: {:#066x}, limit: {:#066x}",
                block_hash, nbits, target, limit
            );
            return Err(BlockError::HashDoesNotMatchNBits);
        }

        self.header.check()?;

        Ok(())
    }

    fn is_all_tx_final(&self, block_time: i64, best_hight: u32) -> bool {
        for tx in &self.transactions {
            if !tx.is_final(block_time, best_hight) {
                return false;
            }
        }
        true
    }

    fn get_serialize_size(&self) -> usize {
        bincode::serialize(self).unwrap().len() // TODO: optimize
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BlockInDb {
    header: BlockHeader,
    transaction_ids: Vec<BlockId>,
}

impl BlockInDb {
    pub fn new(block: &Block) -> Self {
        let txs_num = block.block.transactions.len();
        Self {
            header: block.block.header.clone(),
            transaction_ids: block.merkle_tree[0..txs_num].to_vec(),
        }
    }

    pub fn dump(&self) {
        println!("BlockInDb Data:");
        self.header.dump();
        println!("  Transaction Ids:");
        for id in &self.transaction_ids {
            println!("    {}", id.to_hex_string());
        }
    }

    pub fn calc_block_id(&self) -> BlockId {
        self.header.calc_block_id()
    }

    pub fn get_block_header(&self) -> &BlockHeader {
        &self.header
    }

    pub fn get_transaction_ids(&self) -> &Vec<BlockId> {
        &self.transaction_ids
    }

    pub fn get_prev_block_id(&self) -> &BlockId {
        &self.header.prev_block_header_hash
    }

    pub fn get_version(&self) -> u32 {
        self.header.version
    }

    pub fn get_timestamp(&self) -> u32 {
        self.header.timestamp
    }
    pub fn get_nbits(&self) -> u32 {
        self.header.nbits
    }
    pub fn get_nonce(&self) -> u32 {
        self.header.nonce
    }
    pub fn get_merkle_root_hash(&self) -> &BlockChainHash {
        &self.header.merkle_root_hash
    }
}

//
// Serialized Blocks
//      https://developer.bitcoin.org/reference/block_chain.html
//
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Block {
    block: BlockPayload,
    merkle_tree: Vec<BlockChainHash>,
}

impl Block {
    pub fn new(header: BlockHeader) -> Self {
        Self {
            block: BlockPayload::new(header),
            merkle_tree: vec![],
        }
    }

    pub fn new_null() -> Self {
        Self {
            block: BlockPayload::new(BlockHeader::new_null()),
            merkle_tree: vec![],
        }
    }

    pub fn new_with_nbits(nbits: u32, prev_block_header_hash: &BlockId) -> Self {
        Self {
            block: BlockPayload::new(BlockHeader::new_with_nbits_prev_block_id(
                nbits,
                prev_block_header_hash,
            )),
            merkle_tree: vec![],
        }
    }

    pub fn dump(&self) {
        println!("Block Data:");
        self.block.dump();
        println!("  Merkle Tree:");
        for hash in &self.merkle_tree {
            println!("    {}\n", hash.to_hex_string());
        }
    }

    pub fn preliminary_check(&self) -> Result<(), BlockError> {
        self.block.check()?;

        let merkle_root = self.block.header.get_merkle_root();
        let calc_merkle_root = self.calc_merkle_tree_root()?;
        if calc_merkle_root != merkle_root {
            error!(
                "[Merkle root mismatch] calc_merkle_root({}), merkle_root({})",
                calc_merkle_root, merkle_root
            );
            return Err(BlockError::HashMerkleRootMismatch);
        }

        Ok(())
    }

    pub fn calc_block_id(&self) -> BlockId {
        self.block.header.calc_block_id()
    }

    pub fn add_transaction(&mut self, tx: Transaction) {
        self.block.add_transaction(tx);
    }

    pub fn delete_transaction(&mut self, tx_id: &TransactionId) -> usize {
        self.block.delete_transaction(tx_id)
    }

    pub fn get_transactions(&self) -> &Vec<Transaction> {
        self.block.get_transactions()
    }

    pub fn get_merkle_tree(&self) -> &Vec<BlockId> {
        &self.merkle_tree
    }

    pub fn get_prev_block_id(&self) -> &BlockChainHash {
        &self.block.header.prev_block_header_hash
    }

    pub fn get_timestamp(&self) -> u32 {
        self.block.header.timestamp
    }

    pub fn set_timestamp(&mut self, block_time: u32) {
        self.block.header.timestamp = block_time;
    }

    pub fn get_merkle_root_hash(&self) -> BlockChainHash {
        self.block.header.merkle_root_hash
    }

    pub fn get_version(&self) -> u32 {
        self.block.header.version
    }

    pub fn get_nonce(&self) -> u32 {
        self.block.header.nonce
    }

    // Ok(merkle_tree_root_hash)
    pub fn build_merkle_tree(&mut self) -> Result<BlockChainHash, BlockError> {
        let mut merkle_tree: Vec<BlockId> = vec![];
        let root = self.do_build_merkle_tree(&mut merkle_tree)?;
        self.merkle_tree = merkle_tree;

        Ok(root)
    }

    fn calc_merkle_tree_root(&self) -> Result<BlockChainHash, BlockError> {
        let mut merkle_tree: Vec<BlockId> = vec![];
        let root = self.do_build_merkle_tree(&mut merkle_tree)?;

        Ok(root)
    }

    fn do_build_merkle_tree(
        &self,
        merkle_tree: &mut Vec<BlockId>,
    ) -> Result<BlockChainHash, BlockError> {
        if self.block.transactions.is_empty() {
            return Err(BlockError::NoTransaction);
        }

        let mut current_layer: Vec<BlockChainHash> = self
            .block
            .transactions
            .iter()
            .map(|tx| {
                let tx_raw = bincode::serialize(tx).unwrap();
                let hash = Sha256::digest(&tx_raw);
                let hash_bytes: [u8; 32] = hash.into();
                BlockChainHash::new(hash_bytes)
            })
            .collect();

        // let mut merkle_tree: Vec<BlockId> = current_layer.clone();
        merkle_tree.extend(current_layer.clone());

        while current_layer.len() > 1 {
            if current_layer.len() % 2 != 0 {
                current_layer.push(*current_layer.last().unwrap());
            }

            let mut next_layer: Vec<BlockChainHash> = Vec::new();

            for chunk in current_layer.chunks(2) {
                let mut hasher = Sha256::new();
                let chunk_raw_0 = bincode::serialize(&chunk[0]).unwrap();
                let chunk_raw_1 = bincode::serialize(&chunk[1]).unwrap();
                hasher.update(&chunk_raw_0);
                hasher.update(&chunk_raw_1);
                next_layer.push(BlockChainHash::new(
                    hasher.finalize().to_vec().try_into().unwrap(),
                ));
            }

            merkle_tree.extend(next_layer.clone());
            current_layer = next_layer;
        }

        Ok(*merkle_tree.last().unwrap())
    }

    pub fn set_merkle_root(&mut self, merkle_root_hash: BlockChainHash) {
        self.block.header.set_merkle_root(merkle_root_hash);
    }

    pub fn find_and_set_valid_nonce(&mut self) -> BlockId {
        let target = self.block.header.nbits_to_hash();
        self.block.header.find_and_set_valid_nonce(&target)
    }

    pub fn is_all_tx_final(&self, block_time: i64, best_hight: u32) -> bool {
        self.block.is_all_tx_final(block_time, best_hight)
    }

    pub fn get_nbits(&self) -> u32 {
        self.block.header.nbits
    }

    // serialize
    pub fn serialize(&self) -> Vec<u8> {
        bincode::serialize(&self.block).unwrap()
    }

    // deserialize
    pub fn deserialize(bytes: &[u8]) -> Result<Self, BlockError> {
        if let Ok(internal_block) = bincode::deserialize::<BlockPayload>(bytes) {
            let merkle_tree: Vec<BlockChainHash> = vec![];
            let mut block = Self {
                block: internal_block,
                merkle_tree,
            };
            let merkle_root_hash = block.build_merkle_tree()?;
            if merkle_root_hash != block.block.header.merkle_root_hash {
                return Err(BlockError::MerkleRootHashError);
            }

            Ok(block)
        } else {
            Err(BlockError::DeserializeError)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::{
        CompactSignatureBytes, Hash256Bytes, OpCode, PublicKeyBytes, Script, TransactionId,
        TransactionValue, block,
        constants::{
            BLOCK_CHAIN_HASH_BYTES_LENGTH, BLOCK_CHIAN_VERSION, PROOF_OF_WORK_LIMIT,
            PROOF_OF_WORK_LIMIT_NBITS, PROOF_OF_WORK_LIMIT_VALID_BYTES_NUM,
        },
        transaction::{Transaction, TransactionIn, TransactionOut, TransactionOutPoint},
    };

    #[test]
    fn test_block_header_len() {
        assert_eq!(80, std::mem::size_of::<BlockHeader>());
    }

    #[test]
    fn test_nbits_to_target() {
        let mut header = BlockHeader::new(
            1,
            &BlockChainHash::new([0; 32]),
            BlockChainHash::new([0; 32]),
            0,
            0,
            0,
        );

        // Example nBits = 0x1d00ffff; // Difficulty 1
        let buf = b"00000000FFFF0000000000000000000000000000000000000000000000000000";
        test_nbits_target(0x1d00ffff, &mut header, buf);

        // Another example: nBits = 0x1b0404cb
        let buf = b"00000000000404CB000000000000000000000000000000000000000000000000";
        test_nbits_target(0x1b0404cb, &mut header, buf);

        // Another example: nBits = 0x01003456
        let buf = b"0000000000000000000000000000000000000000000000000000000000000000";
        test_nbits_target(0x01003456, &mut header, buf);

        // Another example: nBits = 0x01123456
        let buf = b"0000000000000000000000000000000000000000000000000000000000000012";
        test_nbits_target(0x01123456, &mut header, buf);

        // nBits = 0x02008000
        let buf = b"0000000000000000000000000000000000000000000000000000000000000080";
        test_nbits_target(0x02008000, &mut header, buf);

        // nBits = 0x05009234
        let buf = b"0000000000000000000000000000000000000000000000000000000092340000";
        test_nbits_target(0x05009234, &mut header, buf);

        // nBits = 0x04923456
        let buf = b"0000000000000000000000000000000000000000000000000000000092345600";
        test_nbits_target(0x04923456, &mut header, buf);

        // nBits = 0x04123456
        let buf = b"0000000000000000000000000000000000000000000000000000000012345600";
        test_nbits_target(0x04123456, &mut header, buf);
    }

    fn test_nbits_target(nbits: u32, header: &mut BlockHeader, buf: &[u8]) {
        header.nbits = nbits;
        let target = header.nbits_to_big_uint();

        let expected_target = BigUint::parse_bytes(buf, 16).unwrap();

        assert_eq!(target, expected_target);
    }

    #[test]
    fn test_proof_of_work_limit_const() {
        let nbits = PROOF_OF_WORK_LIMIT_NBITS;
        let target = block::nbits_to_hash(nbits);
        println!("target: {}", target);
        let cnt = PROOF_OF_WORK_LIMIT_VALID_BYTES_NUM;
        assert_eq!(target.as_bytes()[..cnt], PROOF_OF_WORK_LIMIT[..cnt]);
    }

    #[test]
    fn test_nbits_to_target_hash() {
        let mut header = BlockHeader::new(
            1,
            &BlockChainHash::new([0; 32]),
            BlockChainHash::new([0; 32]),
            0,
            0,
            0,
        );

        // Example nBits = 0x1d00ffff;
        let buf = "00000000FFFF0000000000000000000000000000000000000000000000000000";
        let hex = hex::decode(buf).unwrap();
        assert_eq!(hex.len(), 32);
        let hex: [u8; 32] = hex.try_into().unwrap();
        let target_hash = BlockChainHash::new(hex);

        header.nbits = 0x1d00ffff;
        let hash = header.nbits_to_hash();

        assert_eq!(hash, target_hash);
    }

    fn create_tx_input_script(sig: CompactSignatureBytes, pub_key: PublicKeyBytes) -> Script {
        let mut script: Script = Script::new_empty();
        script.push(OpCode::OpSignature(sig));
        script.push(OpCode::OpPubKey(pub_key));
        script
    }

    fn create_tx_output_script(pub_key: PublicKeyBytes) -> Script {
        let mut script: Script = Script::new_empty();
        script.push(OpCode::OpDup);
        script.push(OpCode::OpHash160);
        script.push(OpCode::OpPubKeyHash(pub_key.to_public_key_hash()));
        script.push(OpCode::OpEqualverify);
        script.push(OpCode::OpChecksig);
        script
    }

    fn new_tx_out_point(tx_id: [u8; 32], index: u32) -> TransactionOutPoint {
        let tx_id = TransactionId::new(tx_id);
        TransactionOutPoint::new(tx_id, index)
    }

    fn new_tx_out(value: TransactionValue) -> TransactionOut {
        let pub_key: PublicKeyBytes = PublicKeyBytes::new([0; 33]);
        let script_pubkey = create_tx_output_script(pub_key);
        TransactionOut::new(value, script_pubkey)
    }

    fn new_tx_in(sequence: u32) -> TransactionIn {
        let out_point = new_tx_out_point([0; 32], 0);
        let sig = CompactSignatureBytes::new([0; 64]);
        let pub_key: PublicKeyBytes = PublicKeyBytes::new([0; 33]);

        let script_signature = create_tx_input_script(sig, pub_key);
        TransactionIn::new_with_sequence(out_point, script_signature, sequence)
    }

    fn new_tx(sequence: u32, value: TransactionValue) -> Transaction {
        let inputs: Vec<TransactionIn> = vec![new_tx_in(sequence)];
        let outputs: Vec<TransactionOut> = vec![new_tx_out(value)];
        Transaction::new(BLOCK_CHIAN_VERSION, inputs, outputs, 0)
    }

    fn calc_merkle_hash(tx: &Transaction) -> BlockChainHash {
        let tx_raw = bincode::serialize(tx).unwrap();
        let hash = Sha256::digest(&tx_raw);
        let hash_bytes: [u8; 32] = hash.into();
        BlockChainHash::new(hash_bytes)
    }

    #[test]
    fn test_merkle_tree() {
        let header = BlockHeader::new(
            1,
            &BlockChainHash::new([0; 32]),
            BlockChainHash::new([0; 32]),
            0,
            0,
            0,
        );

        let mut block = Block::new(header);

        let mut sequence: u32 = 0xffff_ffff;
        let mut value: TransactionValue = 0;
        let tx0 = new_tx(sequence, value);
        // let tx0_id = tx0.calc_tx_id();
        let tx0_merkle_hash = calc_merkle_hash(&tx0);
        block.block.transactions.push(tx0);

        sequence -= 1;
        value += 1;
        let tx1 = new_tx(sequence, value);
        let tx1_merkle_hash = calc_merkle_hash(&tx1);
        // let tx1_id = tx1.calc_tx_id();
        block.block.transactions.push(tx1);

        let merkle_tree = block.build_merkle_tree();
        assert!(merkle_tree.is_ok());
        assert_eq!(block.merkle_tree.len(), 3);
        assert_eq!(
            block.merkle_tree[0].as_bytes_ref(),
            tx0_merkle_hash.as_bytes_ref()
        );
        assert_eq!(
            block.merkle_tree[1].as_bytes_ref(),
            tx1_merkle_hash.as_bytes_ref()
        );

        let mut hasher = Sha256::new();
        hasher.update(tx0_merkle_hash.as_bytes_ref());
        hasher.update(tx1_merkle_hash.as_bytes_ref());
        let hash: Hash256Bytes = hasher.finalize().to_vec().try_into().unwrap();

        assert_eq!(block.merkle_tree[2].as_bytes_ref(), &hash);

        // Ensure the root of the Merkle tree is computed correctly
        let merkle_tree_root_hash = merkle_tree.unwrap();
        println!("Merkle tree root hash: {:x?}", merkle_tree_root_hash);
    }

    #[test]
    fn test_header_nonce() {
        let version = 1;
        let prev_block_hash = BlockChainHash::new([0; BLOCK_CHAIN_HASH_BYTES_LENGTH]);
        let merkle_root_hash = BlockChainHash::new([1; BLOCK_CHAIN_HASH_BYTES_LENGTH]);
        let timestamp = 1234567890;
        let nbits = PROOF_OF_WORK_LIMIT_NBITS;
        let nonce = 0;
        let mut h = BlockHeader::new(
            version,
            &prev_block_hash,
            merkle_root_hash,
            timestamp,
            nbits,
            nonce,
        );

        let target = h.nbits_to_hash();
        let hash = h.find_and_set_valid_nonce(&target);

        println!("hash: {}", hash.to_hex_string());
        println!("header: {:?}", h);
        assert_eq!(h.calc_block_id(), hash);
    }

    #[test]
    fn test_using_flow_of_block() {
        let header = block::create_block_header_for_mining(
            &BlockChainHash::new([1; BLOCK_CHAIN_HASH_BYTES_LENGTH]),
            0x1eff_ffff,
        );
        let mut block = Block::new(header);

        // let tx_id = TransactionId::new([2; 32]);
        let sequence: u32 = 0xffff_ffff;
        let value: TransactionValue = 0;
        let tx0 = new_tx(sequence, value);
        block.add_transaction(tx0.clone());

        assert_eq!(block.block.transactions.len(), 1);
        assert_eq!(block.block.transactions[0], tx0);

        assert_eq!(block.merkle_tree.len(), 0);
        let _merkle_root_hash = block.build_merkle_tree().unwrap();
        assert_eq!(block.merkle_tree.len(), 1);

        // let tx_id = TransactionId::new([3; 32]);
        let sequence: u32 = 0xffff_fff1;
        let value: TransactionValue = 1;
        let tx1 = new_tx(sequence, value);
        block.add_transaction(tx1);
        let merkle_root_hash = block.build_merkle_tree().unwrap();
        assert_eq!(block.merkle_tree.len(), 3);

        block.set_merkle_root(merkle_root_hash);
        assert_eq!(block.block.header.merkle_root_hash, merkle_root_hash);

        let _block_hash = block.find_and_set_valid_nonce();
    }
}
